package gsthrottle

import (
	"context"
	"sync"
	"time"
)

// 滑动窗口的bucket最大数量
const WINDOW_BUCKET_SIZE_MAX = 60

// 滑动窗口tick时钟跳动的最低周期时间
const WINDOW_TICKMS_MIN = time.Millisecond * 50

// 滑动窗口限流
type WindowThrottle struct {
	tickMs  time.Duration  // 格子周期
	tick    map[int64]uint // 格子位置
	limit   map[int64]uint // 限制数量
	limitMs time.Duration  // 限制时间频率单位
	bucket  map[int64][]uint
	sync.Mutex
	ctx        context.Context
	cancel     context.CancelFunc
	bucketSize int
}

func (x *WindowThrottle) AddThrottle(throttleConfig ...*ThrottleConfig) {
	x.Lock()
	defer x.Unlock()
	for _, v := range throttleConfig {
		x.bucket[v.Id] = make([]uint, x.bucketSize)
		x.limit[v.Id] = v.Limit
	}
}

func (x *WindowThrottle) UpdateThrottle(throttleConfig *ThrottleConfig) {
	x.Lock()
	defer x.Unlock()
	if _, ok := x.bucket[throttleConfig.Id]; ok {
		x.bucket[throttleConfig.Id] = make([]uint, x.bucketSize)
		x.limit[throttleConfig.Id] = throttleConfig.Limit
	}
}

func (x *WindowThrottle) RemoveId(id int64) bool {
	x.Lock()
	defer x.Unlock()
	if _, ok := x.bucket[id]; ok {
		delete(x.bucket, id)
		return true
	}

	return false
}

func (x *WindowThrottle) Init() {
	// bucket不能太大
	size := x.limitMs / x.tickMs
	if size > WINDOW_BUCKET_SIZE_MAX {
		size = WINDOW_BUCKET_SIZE_MAX
		x.tickMs = x.limitMs / size
	}

	// tickms时钟不能太小
	if x.tickMs < WINDOW_TICKMS_MIN {
		x.tickMs = WINDOW_TICKMS_MIN
		size = x.limitMs / x.tickMs
	}

	x.ctx, x.cancel = context.WithCancel(context.Background())
	x.bucketSize = int(size)
	x.bucket = make(map[int64][]uint)
	ticker := time.NewTicker(x.tickMs)

	go func() {
		for {
			select {
			case <-x.ctx.Done():
				return
			case <-ticker.C:
				x.Lock()
				x.Unlock()
				for k := range x.bucket {
					x.tickNext(k)
				}
			}
		}
	}()
}

func (x *WindowThrottle) tickNext(id int64) {
	x.tick[id]++
	if x.isBucketTail(id) {
		x.tick[id] = 0
	}

	x.bucket[id][x.tick[id]] = 0
}

func (x *WindowThrottle) sumBucket(id int64) uint {
	res := uint(0)
	for _, v := range x.bucket[id] {
		res = res + v
	}

	return res
}

func (x *WindowThrottle) isBucketTail(id int64) bool {
	return x.tick[id] >= uint(len(x.bucket[id]))
}

func (x *WindowThrottle) isBucketFull(id int64) bool {
	return x.sumBucket(id) >= x.limit[id]
}

func (x *WindowThrottle) bucketInc(id int64) {
	x.bucket[id][x.tick[id]]++
}

func (x *WindowThrottle) Call(id int64, call func(), reject func()) {
	x.Lock()
	defer x.Unlock()
	if x.isBucketFull(id) {
		reject()
		return
	}

	x.bucketInc(id)
	call()
}
